import os.path

import uiautomator2 as u2
import numpy as np
import cv2
from PIL import Image
from io import BytesIO
import time

from utils.logger import log
from utils.utils import CommonUtils


def timestamped_filename(prefix="screenshot", ext=".png"):
    return f"{prefix}_{time.strftime('%Y%m%d_%H%M%S')}{ext}"


class ImageMatcher:
    def __init__(self, client=None, ip=None, port=None, pk=None):
        self.client = client or u2.connect(f"{ip}:{port}")
        self.pk = pk

    def save_template(self, x, y, w, h, name):
        pil_img = self.client.screenshot(format='pillow')
        copped_img = pil_img.crop((x, y, x+w, y+h))
        shot_path = CommonUtils.resource_path(["resource", f"{name}.png"])
        if os.path.exists(shot_path):
            return False, shot_path
        copped_img.convert("L").save(shot_path)
        return True, shot_path

    def screenshot_cv2(self):
        """获取手机截图并转为 OpenCV 格式"""
        time.sleep(1)
        pil_img = self.client.screenshot(format='pillow')
        cv_img = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2GRAY)

        return cv_img

    def load_template(self, path_or_bytes):
        """加载模板图像为 OpenCV 格式"""
        if isinstance(path_or_bytes, str):
            if not os.path.exists(path_or_bytes):
                raise FileNotFoundError(f"{path_or_bytes}")
            return cv2.imread(path_or_bytes, cv2.IMREAD_GRAYSCALE)
        elif isinstance(path_or_bytes, bytes):
            image = Image.open(BytesIO(path_or_bytes)).convert("L")
            return np.array(image)
        else:
            raise TypeError("模板类型必须是路径字符串或字节流")

    def _find_image(self, template_cv, threshold=0.8, all=False):
        """返回第一个匹配位置(center_x, center_y)，如果没有匹配，返回 None"""
        screen = self.screenshot_cv2()
        res = cv2.matchTemplate(screen, template_cv, cv2.TM_CCOEFF_NORMED)
        loc = np.where(res >= threshold)
        h, w = template_cv.shape[:2]
        if all:
            all_list = []
            for pt in zip(*loc[::-1]):
                center_x = pt[0] + w // 2
                center_y = pt[1] + h // 2
                all_list.append([float(center_x), float(center_y)])
            return all_list
        else:
            for pt in zip(*loc[::-1]):
                center_x = pt[0] + w // 2
                center_y = pt[1] + h // 2
                return float(center_x), float(center_y)
        return None

    def click_image(self, template, threshold=0.8, max_retries=1, interval=1.0):
        """
        查找并点击图像位置
        :param template: 模板路径或字节流
        :param threshold: 相似度阈值
        :param max_retries: 重试次数
        :param interval: 每次重试间隔秒数
        :return: 是否点击成功
        """
        template_cv = self.load_template(template)

        for attempt in range(max_retries):
            pos = self._find_image(template_cv, threshold)
            if pos:
                log.info(f"[✔]识别到:{os.path.split(template)[-1].split('.')[0]}", extra={"pk": self.pk})
                self.client.click(*pos)
                return True
            else:
                log.info(f"[ ]未识别到:{os.path.split(template)[-1].split('.')[0]} -> {attempt+1}次", extra={"pk": self.pk})
        return False

    def find_image(self, template, threshold=0.8, max_retries=1, interval=1.0):
        """
        查找并点击图像位置
        :param template: 模板路径或字节流
        :param threshold: 相似度阈值
        :param max_retries: 重试次数
        :param interval: 每次重试间隔秒数
        :return: 是否点击成功
        """
        template_cv = self.load_template(template)

        for attempt in range(max_retries):
            pos = self._find_image(template_cv, threshold)
            if pos:
                log.info(f"[✔]识别到:{os.path.split(template)[-1].split('.')[0]}", extra={"pk": self.pk})
                return True
            else:
                log.info(f"[ ]未识别到:{os.path.split(template)[-1].split('.')[0]} -> {attempt+1}次", extra={"pk": self.pk})
        return False

    def find_all_image(self, template, threshold=0.8, max_retries=1, interval=1.0):
        """
        查找并点击图像位置
        :param template: 模板路径或字节流
        :param threshold: 相似度阈值
        :param max_retries: 重试次数
        :param interval: 每次重试间隔秒数
        :return: 是否点击成功
        """
        template_cv = self.load_template(template)

        pos = []
        for attempt in range(max_retries):
            pos = self._find_image(template_cv, threshold, all=True)
            if pos:
                break

        return pos
